from lark import Lark, Transformer
from dataclasses import dataclass
from typing import Union, List


@dataclass
class BasicType:
    """C++基本类型"""

    type: str


@dataclass
class CustomizedType:
    """用户自定义类型"""

    type: str


@dataclass
class MapType:
    """Map类型"""

    # 键值类型应该更加严格，MapType应该不适用
    key_type: Union[BasicType, CustomizedType, "VectorType", "MapType", "ArrayType"]
    value_type: Union[BasicType, CustomizedType, "VectorType", "MapType", "ArrayType"]

@dataclass
class ArrayType:
    type: Union[BasicType, CustomizedType, "VectorType", "MapType", "ArrayType"]
    num: int

@dataclass
class VectorType:
    """Vector类型"""

    type: Union[BasicType, CustomizedType, "VectorType", "MapType", "ArrayType"]


@dataclass
class MemberDecl:
    """类成员声明"""

    type: List[Union[BasicType, CustomizedType, "VectorType", "MapType", "ArrayType"]]
    name: str


@dataclass
class ClassType:
    """class类型"""

    name: str
    fields: List[MemberDecl]
    inner_classes: List["ClassType"]


@dataclass
class StructType:
    """struct类型"""

    name: str
    fields: List[MemberDecl]


@dataclass
class EnumItem:
    name: str
    val: int


@dataclass
class EnumType:
    name: str
    underlying_type: BasicType
    enumItems: List[EnumItem]
    value_count: int


class DSLTransformer(Transformer):
    def __init__(self):
        self.messages = []

    def start(self, items):
        return items

    def message(self, items):
        return items[0]

    def struct_decl(self, items):
        return StructType(items[0], items[1:])

    def class_decl(self, items):
        name = items[0]
        fields = []
        inner_classes = []
        for item in items[1:]:
            if isinstance(item, MemberDecl):
                fields.append(item)
            if isinstance(item, ClassType):
                inner_classes.append(item)
            if isinstance(item, StructType):
                inner_classes.append(item)
            if isinstance(item, EnumType):
                inner_classes.append(item)
        if len(fields) == 0:
            raise ValueError(f"class {name} must have at least one field.")
        return ClassType(name, fields, inner_classes)

    def member_decl(self, items):
        return MemberDecl(type=items[0], name=items[1])

    def container_decl(self, items):
        return items[0]

    def vector_decl(self, items):
        return VectorType(items[0])
    
    def array_decl(self, items):
        return ArrayType(type=items[0], num=int(items[1]))

    def map_decl(self, items):
        key_type, value_type = items
        return MapType(key_type, value_type)

    def type_decl(self, items):
        return items[0]

    def BASIC_TYPE(self, t):
        return BasicType(t.value)

    def custom_type(self, parts):
        return CustomizedType(".".join(p for p in parts))

    def CNAME(self, token):
        return token.value

    def enum_decl(self, items):
        # 这里可以做一些数值和enum底层表示类型等等的检查
        return EnumType(items[0], items[1], items[2:], len(items[2:]))

    def enum_item_decl(self, items):
        return EnumItem(name=items[0], val=int(items[1].value))


def parse_dsl(schema_content):
    with open("dsl/dsl.lark") as f:
        grammar = f.read()
    parser = Lark(grammar, parser="lalr", transformer=DSLTransformer())
    result = parser.parse(schema_content)
    if isinstance(result, list):
        return result
    return [result]
